import logging
import os
import pickle
import uuid
from enum import Enum
from typing import Any, List, Optional, Tuple

import redis.asyncio as async_redis
from pydantic import BaseModel

from .dataclasses.common import Message
from .dataclasses.knowledge import Graph


class Role(Enum):
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"


class SessionState(Enum):
    WAITING = "waiting"  # 等待用户输入
    PROCESSING = "processing"  # 正在处理
    FINISHED = "finished"  # 处理完成


class Message(BaseModel):
    role: Role  # Changed from str to Role enum
    content: str
    name: Optional[str] = None

    def to_str(self):
        return f"{self.role.value}: {self.content}"


class Status(str, Enum):
    REASK_LEVEL_1 = "reask_level_1"  # 反问一级事项
    REASK_LEVEL_2 = "reask_level_2"  # 反问二级事项
    REASK_SCENE = "reask_scene"  # 反问场景
    REASK_INFO = "reask_info"  # 反问信息，找工作或者计算场景
    INFO_ANSWER = "info_answer"  # 信息回答
    RAG_ANSWER = "answer"  # RAG回答
    GRAPH_ANSWER = "graph_answer"  # 图谱回答
    REFUSE = "refuse"  # 拒答


class Session(BaseModel):
    session_id: str = ""
    history: List[Message] = []
    graph: Optional[Graph] = None
    categories: List[str] = []
    extra: dict = {}  # 添加缺失的 extra 字段

    def add_history(self, role: Role, content: str):
        """Add a message to the history."""
        if self.history and self.history[-1].role == role:  # Compare Role enums
            raise ValueError("Cannot add history with the same role back-to-back.")
        self.history.append(Message(role=role, content=content))

    def get_message_dict(self):
        """Get messages in dict format."""
        return [
            {
                "role": msg.role.value,  # Access enum value
                "content": msg.content,
            }
            for msg in self.history
        ]

    def is_empty(self):
        return len(self.history) == 0

    def put_extra(self, key: str, value: Any):
        self.extra[key] = value

    def cur_turn(self) -> int:
        """Count number of USER messages."""
        return len([msg for msg in self.history if msg.role == Role.USER])

    def validate(self) -> bool:
        if len(self.history) == 0 or self.history[-1].role == Role.ASSISTANT:
            return True
        else:
            return False

    def fix(self):
        if self.validate():
            return

        self.history.pop()

    def update_graph(self, graph: Graph, **kwargs):
        if not self.graph:
            self.graph = graph
        else:
            self.graph.update(graph, **kwargs)

    def latest(self, role: Role = None) -> str:  # 修改返回类型为 str
        if not self.history:
            return ""

        if role is None:
            return self.history[-1].content
        else:
            for msg in reversed(self.history):
                if msg.role == role:
                    return msg.content
            return ""

    def user_history(self):
        return [msg.content for msg in self.history if msg.role == Role.USER]

    def chat_history(self):
        text = ""
        for msg in self.history:
            text += f"{msg.role.value}: {msg.content}\n"
        return text


class SessionManager:
    def __init__(
        self,
        redis_host: str = "redis://localhost:6379/0",
        timeout=120,
        backup_timeout=240,  # 备份保留更长时间
        password=None,
        logger=None,
    ):
        self.redis_host = redis_host
        self.redis = async_redis.from_url(self.redis_host, password=password, db=1)
        self.timeout = timeout
        self.backup_timeout = backup_timeout
        self.logger = logger or logging.getLogger(__name__)

    @classmethod
    def from_env(cls):
        redis_host = os.getenv("REDIS_HOST", "redis://localhost:6379/0")
        timeout = int(os.getenv("SESSION_TIMEOUT", 120))
        backup_timeout = int(
            os.getenv("BACKUP_TIMEOUT", timeout * 2)
        )  # 备份超时时间默认为普通超时的两倍
        password = os.getenv("REDIS_PASSWORD", None)

        return cls(
            redis_host=redis_host,
            timeout=timeout,
            backup_timeout=backup_timeout,
            password=password,
        )

    async def _get_state(
        self, session_id: str
    ) -> Tuple[Optional[str], Optional[bytes]]:
        """获取 session 状态和数据，处理错误"""
        try:
            # 使用 pipeline 减少网络往返
            pipeline = self.redis.pipeline()
            pipeline.get(session_id + "_state")
            pipeline.get(session_id)
            results = await pipeline.execute()

            state_bytes, session_bytes = results
            state = state_bytes.decode() if state_bytes else None

            return state, session_bytes
        except Exception as e:
            self.logger.error(f"Error getting session state: {e}")
            return None, None

    async def put(self, session: Session):
        """存储 session 并设置状态"""
        try:
            # 使用 pipeline 确保原子性
            pipeline = self.redis.pipeline()
            serialized_session = pickle.dumps(session)
            pipeline.setex(session.session_id, self.timeout, serialized_session)
            pipeline.setex(
                session.session_id + "_state", self.timeout, SessionState.WAITING.value
            )
            await pipeline.execute()
            return True
        except Exception as e:
            self.logger.error(f"Error saving session: {e}")
            return False

    async def mark_finished(self, session_id: str):
        """将 session 标记为已完成"""
        try:
            await self.redis.setex(
                session_id + "_state", self.timeout, SessionState.FINISHED.value
            )
            return True
        except Exception as e:
            self.logger.error(f"Error marking session as finished: {e}")
            return False

    async def pop(self, session_id: str, **kwargs) -> Session:
        """获取 session 并标记为正在处理"""
        try:
            # 获取当前状态和 session 数据
            cur_state, state_data = await self._get_state(session_id)

            # 创建新 session 的情况
            if state_data is None and cur_state is None:
                await self.redis.setex(
                    session_id + "_state", self.timeout, SessionState.PROCESSING.value
                )
                return self.new_session(session_id, **kwargs)

            # session 正在被处理，尝试从备份恢复
            elif cur_state == SessionState.PROCESSING.value:
                backup_data = await self.redis.get(session_id + "_backup")
                if backup_data is None:
                    raise RuntimeError(
                        "Session is being processed and no backup found."
                    )

                await self.redis.expire(
                    session_id + "_state", self.timeout
                )  # 刷新状态超时
                session = pickle.loads(backup_data)
                return session

            # session 已完成处理
            elif cur_state == SessionState.FINISHED.value:
                raise RuntimeError("Session has been finished.")

            # 正常情况：session 等待处理
            else:
                # 创建事务确保原子性
                transaction = self.redis.pipeline()

                # 1. 先创建备份
                transaction.setex(
                    session_id + "_backup", self.backup_timeout, state_data
                )
                # 2. 设置状态为处理中
                transaction.setex(
                    session_id + "_state", self.timeout, SessionState.PROCESSING.value
                )
                # 3. 删除原始 session（因为我们已经有备份）
                transaction.delete(session_id)

                # 执行事务
                await transaction.execute()

                # 反序列化并返回
                session = pickle.loads(state_data)
                return session

        except Exception as e:
            self.logger.error(f"Error in pop operation: {e}")
            # 尝试恢复备份
            try:
                backup_data = await self.redis.get(session_id + "_backup")
                if backup_data:
                    self.logger.info(f"Recovering session {session_id} from backup")
                    return pickle.loads(backup_data)
            except Exception as recover_error:
                self.logger.error(f"Failed to recover from backup: {recover_error}")

            # 如果都失败，创建新会话
            self.logger.warning(f"Creating new session due to failure")
            return self.new_session(session_id, **kwargs)

    async def get(self, session_id: str) -> Optional[Session]:
        """只读方式获取 session，不改变状态"""
        try:
            state_data = await self.redis.get(session_id)
            if state_data is None:
                return None
            return pickle.loads(state_data)
        except Exception as e:
            self.logger.error(f"Error getting session: {e}")
            return None

    def new_session(self, session_id: Optional[str] = None, **kwargs) -> Session:
        """创建新会话"""
        if session_id is None:
            session_id = str(uuid.uuid4())
        return Session(session_id=session_id, **kwargs)

    async def exists(self, session_id: str) -> bool:
        """检查会话是否存在"""
        try:
            return bool(await self.redis.exists(session_id))
        except Exception as e:
            self.logger.error(f"Error checking session existence: {e}")
            return False

    async def delete(self, session_id: str) -> bool:
        """删除会话及相关状态"""
        try:
            # 使用 pipeline 一次性删除所有相关键
            pipeline = self.redis.pipeline()
            pipeline.delete(session_id)
            pipeline.delete(session_id + "_state")
            pipeline.delete(session_id + "_backup")
            await pipeline.execute()
            return True
        except Exception as e:
            self.logger.error(f"Error deleting session: {e}")
            return False

    async def close(self):
        """正确关闭 Redis 连接"""
        try:
            await self.redis.aclose()
        except Exception as e:
            self.logger.error(f"Error closing Redis connection: {e}")


# 单例模式
session_manager = SessionManager.from_env()
